import os
import sys
from tree_sitter import Language, Parser
from tree_sitter_languages import get_language, get_parser
from typing import Dict, List, Any, Optional, Tuple
import json
import re

class TypeScriptDeclarationsParser:
    """使用 Tree-sitter 解析 TypeScript 声明文件"""

    def __init__(self):
        # 加载编译好的 TypeScript 语法库
        self.language = get_language("typescript")
        self.parser = get_parser("typescript")

        # 定义更全面的查询模式
        self.query_patterns = {
            'function': """
            (function_declaration
                name: (identifier) @function.name
                parameters: (formal_parameters) @function.params
                type: (type_annotation)? @function.return_type
            ) @function.declaration
            """,
            'interface': """
            [
                (export_statement
                    (interface_declaration
                        name: (type_identifier) @interface.name
                        body: (object_type) @interface.body
                    ) @interface.declaration
                )
            ]
            """,
            'type_alias': """
            [
                (type_alias_declaration
                    name: (type_identifier) @type.name
                    value: (function_type
                        parameters: (formal_parameters) @type.params
                        return_type: (type_annotation)? @type.return_type
                    ) @type.value
                ) @type.declaration
                (type_alias_declaration
                    name: (type_identifier) @type.name
                    value: (_) @type.value
                ) @type.declaration
            ]
            """,
            'method_signature': """
            (method_signature
                name: (property_identifier) @method.name
                parameters: (formal_parameters) @method.params
                type: (type_annotation)? @method.return_type
            ) @method.declaration
            """,
             'property_signature': """
            (property_signature
                name: (property_identifier) @property.name
                type: (type_annotation) @property.type
            ) @property.declaration
            """,
            'exported_function': """
            (export_statement
            (lexical_declaration
                (variable_declarator
                    name: (identifier) @function.name
                    type: (type_annotation
                        (function_type
                            parameters: (formal_parameters) @function.params
                            return_type: (_) @function.return_type
                        )
                    )
                )
            )
            ) @export.declaration
            """,
            'exported_const_function': """
            (export_statement
                (variable_declaration
                    (variable_declarator
                        name: (identifier) @function.name
                        value: (arrow_function
                            parameters: (formal_parameters) @function.params
                            return_type: (type_annotation)? @function.return_type
                        )? @function.value
                    )
                )
            ) @export.declaration
            """,
            'exported_variable': """
            (export_statement
                (variable_declaration
                    (variable_declarator
                        name: (identifier) @variable.name
                        type: (type_annotation)? @variable.type
                    )
                )
            ) @export.declaration
            """
        }

        # 编译查询
        self.queries = {}
        for key, pattern in self.query_patterns.items():
            try:
                self.queries[key] = self.language.query(pattern)
            except Exception as e:
                print(f"编译查询 '{key}' 时出错: {e}")
                print(f"查询模式:\n{pattern}")
                self.queries[key] = None

    def parse_file(self, file_path: str) -> Dict[str, Any]:
        """解析 TypeScript 声明文件"""
        with open(file_path, 'r', encoding='utf-8') as f:
            source_code = f.read()

        # 解析代码
        tree = self.parser.parse(bytes(source_code, 'utf-8'))
        # 提取声明
        return self._extract_declarations(tree, source_code)

    def _extract_declarations(self, tree, source_code: str) -> Dict[str, Any]:
        """从语法树中提取声明信息"""
        root_node = tree.root_node

        # 初始化结果字典
        result = {
            'functions': [],
            'interfaces': [],
            'type_aliases': [],
            'classes': [],
            'variables': [],
            'exported_functions': [],
            'exported_variables': []
        }

        # 使用查询来提取各种声明
        self._extract_with_queries(root_node, source_code, result)

        # 额外遍历所有顶级声明作为后备
        for child in root_node.children:
            # 跳过注释
            if child.type in ['comment', 'line_comment', 'block_comment']:
                continue
            # print(f'type: {child.type}')
            if child.type == 'function_declaration':
                func = self._extract_function(child, source_code)
                if func and func['name'] != 'anonymous':
                    result['functions'].append(func)
            elif child.type == 'interface_declaration':
                pass
                # result['interfaces'].append(self._extract_interface(child, source_code))
            elif child.type == 'type_alias_declaration':
                result['type_aliases'].append(self._extract_type_alias(child, source_code))
            elif child.type == 'class_declaration':
                result['classes'].append(self._extract_class(child, source_code))
            elif child.type == 'variable_declaration':
                # 检查是否是导出语句的一部分
                parent = child.parent
                if parent and parent.type == 'export_statement':
                    variables = self._extract_variable(child, source_code)
                    for var_info in variables:
                        result['exported_variables'].append(var_info)
                else:
                    variables = self._extract_variable(child, source_code)
                    for var_info in variables:
                        result['variables'].append(var_info)
            elif child.type == 'export_statement':
                # self.print_syntax_tree(child, source_code)
                pass
        return result
    
    def _find_declaration_root(self, node):
        """查找声明根节点（export_statement 或 lexical_declaration）"""
        while node and node.parent:
            if node.type in ['export_statement', 'lexical_declaration', 'variable_declaration']:
                return node
            node = node.parent
        return node

    def _extract_with_queries(self, node, source_code: str, result: Dict[str, Any]):
        """使用查询模式提取声明"""
        for query_name, query in self.queries.items():
            if query is None:
                continue

            try:
                captures = query.captures(node)
                if query_name == 'exported_function' or query_name == 'exported_const_function':
                    declarations = {}
                    for capture_node, tag in captures:
                        # 获取声明根节点
                        decl_root = self._find_declaration_root(capture_node)
        
                        if decl_root not in declarations:
                            declarations[decl_root] = {}
        
                        declarations[decl_root][tag] = capture_node
    
                    # 处理每个声明
                    for decl_root, tags in declarations.items():
                        if 'function.name' in tags:
                            function_name_node = tags['function.name']
                            function_name = self._get_node_text(function_name_node, source_code)
            
                            # 初始化参数和返回类型
                            params = []
                            return_type = 'void'
            
                            # 获取参数
                            if 'function.params' in tags:
                                params_node = tags['function.params']
                                # 使用改进的参数提取方法
                                params = self._extract_parameters(params_node, source_code)
            
                            # 获取返回类型
                            if 'function.return_type' in tags:
                                return_type_node = tags['function.return_type']
                                return_type = self._get_node_text(return_type_node, source_code)
            
                            # 获取完整声明文本
                            source_text = self._get_node_text(decl_root, source_code)
            
                            # 添加到结果
                            result['exported_functions'].append({
                                'name': function_name,
                                'parameters': params,
                                'return_type': return_type,
                                'source_text': source_text
                            })

                elif query_name == 'exported_variable':
                    for capture in captures:
                        if capture[1] == 'variable.name':
                            var_name = self._get_node_text(capture[0], source_code)
                            var_type = 'any'

                            for cap in captures:
                                if cap[1] == 'variable.type' and self._is_same_declaration(capture[0], cap[0]):
                                    var_type = self._get_node_text(cap[0], source_code)

                            result['exported_variables'].append({
                                'name': var_name,
                                'type': var_type,
                                'source_text': self._get_declaration_text(capture[0], source_code)
                            })

                elif query_name == 'function':
                    for capture in captures:
                        if capture[1] == 'function.name':
                            function_name = self._get_node_text(capture[0], source_code)
                            params = []
                            return_type = 'void'

                            for cap in captures:
                                if cap[1] == 'function.params' and self._is_same_declaration(capture[0], cap[0]):
                                    params = self._extract_parameters(cap[0], source_code)
                                elif cap[1] == 'function.return_type' and self._is_same_declaration(capture[0], cap[0]):
                                    return_type = self._get_node_text(cap[0], source_code)

                            result['functions'].append({
                                'name': function_name,
                                'parameters': params,
                                'return_type': return_type,
                                'source_text': self._get_declaration_text(capture[0], source_code)
                            })
                elif query_name == "interface":
                    for capture_node, tag in captures:
                        if tag == 'interface.declaration':
                            interface_info = self._extract_interface(capture_node, source_code)
                            result['interfaces'].append(interface_info)

            except Exception as e:
                print(f"执行查询 '{query_name}' 时出错: {e}")

    def _extract_interface(self, node, source_code: str) -> Dict[str, Any]:
        """提取接口声明信息"""
        # 如果是导出语句，获取内部的接口声明节点
        if node.type == 'export_statement':
            # 查找接口声明子节点
            for child in node.children:
                if child.type == 'interface_declaration':
                    node = child
                    break
    
        # 获取接口名
        name_node = node.child_by_field_name('name')
        name = self._get_node_text(name_node, source_code) if name_node else 'unknown'
    
        # 获取接口体
        body_node = node.child_by_field_name('body')
        members = self._extract_interface_members(body_node, source_code) if body_node else []
    
        return {
            'name': name,
            'members': members,
            'source_text': self._get_node_text(node, source_code)
        }

    def _is_same_declaration(self, node1, node2) -> bool:
        """检查两个节点是否属于同一个声明"""
        # 查找共同的父节点
        def get_declaration_root(node):
            while node and node.type not in ['function_declaration', 'variable_declarator',
                                           'method_signature', 'property_signature',
                                           'export_statement']:
                node = node.parent
            return node

        return get_declaration_root(node1) == get_declaration_root(node2)

    def _get_declaration_text(self, node, source_code: str) -> str:
        """获取整个声明的文本内容"""
        # 向上查找直到找到完整的声明节点
        while node and node.type not in ['function_declaration', 'variable_declaration',
                                      'interface_declaration', 'type_alias_declaration',
                                      'export_statement']:
            node = node.parent

        return self._get_node_text(node, source_code) if node else ''

    def _extract_function(self, node, source_code: str) -> Dict[str, Any]:
        """提取函数声明信息"""
        # 获取函数名
        name_node = node.child_by_field_name('name')
        name = self._get_node_text(name_node, source_code) if name_node else 'anonymous'

        # 获取参数
        params_node = node.child_by_field_name('parameters')
        parameters = self._extract_parameters(params_node, source_code) if params_node else []

        # 获取返回类型
        return_type_node = node.child_by_field_name('type')
        return_type = self._get_node_text(return_type_node, source_code) if return_type_node else 'void'

        return {
            'name': name,
            'parameters': parameters,
            'return_type': return_type,
            'source_text': self._get_node_text(node, source_code)
        }

    def _extract_exported_functions(self, node, source_code: str) -> Dict[str, Any]:
        """提取函数声明信息"""
        # 获取函数名
        name_node = node.child_by_field_name('name')
        name = self._get_node_text(name_node, source_code) if name_node else 'anonymous'

        value_node = node.child_by_field_name('type')
        print(f'{value_node},    {node}')
        # 获取参数
        params_node = value_node.child_by_field_name('parameters')
        parameters = self._extract_parameters(params_node, source_code) if params_node else []

        # 获取返回类型
        return_type_node = value_node.child_by_field_name('return_type')
        return_type = self._get_node_text(return_type_node, source_code) if return_type_node else 'void'

        return {
            'name': name,
            'parameters': parameters,
            'return_type': return_type,
            'source_text': self._get_node_text(node, source_code)
        }
    
    def _is_optional_property(self, node) -> bool:
        """检查属性是否可选"""
        # 1. 检查节点类型是否为可选参数
        if node.type == 'optional_parameter':
            return True
    
        # 2. 检查是否有问号标记
        for child in node.children:
            if child.type == 'question_mark':
                return True
        return False
    
    def _extract_property_type(self, property_node, source_code: str) -> str:
        """提取属性类型文本（不包括冒号）"""
        # 查找类型注解节点
        type_node = property_node.child_by_field_name('type')
        if not type_node:
            return 'any'
    
        # 获取完整类型注解文本
        full_text = self._get_node_text(type_node, source_code)
    
        # 移除开头的冒号和空格
        if full_text.startswith(':'):
            # 移除冒号及后续空格
            return full_text[1:].lstrip()
    
        # 处理可选属性后紧跟冒号的情况
        if full_text.startswith('?:') or full_text.startswith('? :'):
            # 移除问号、冒号及后续空格
            return full_text.split(':', 1)[1].lstrip()
    
        return full_text

    def _extract_interface_members(self, node, source_code: str) -> List[Dict[str, Any]]:
        """提取接口成员（属性和方法）"""
        members = []

        for child in node.children:
            # 跳过注释
            if child.type in ['comment', 'line_comment', 'block_comment']:
                continue

            if child.type == 'property_signature':
                # 提取属性
                name_node = child.child_by_field_name('name')
                name = self._get_node_text(name_node, source_code) if name_node else 'unknown'

                type_node = child.child_by_field_name('type')
                # type_text = self._get_node_text(type_node, source_code) if type_node else 'any'
                type_text = self._extract_property_type(child, source_code)
                # 检查是否可选
                optional = self._is_optional_property(child)

                members.append({
                    'type': 'property',
                    'name': name,
                    'data_type': type_text,
                    'optional': optional
                })

            elif child.type == 'method_signature':
                # 提取方法
                name_node = child.child_by_field_name('name')
                name = self._get_node_text(name_node, source_code) if name_node else 'unknown'

                params_node = child.child_by_field_name('parameters')
                parameters = self._extract_parameters(params_node, source_code) if params_node else []

                return_type_node = child.child_by_field_name('type')
                return_type = self._get_node_text(return_type_node, source_code) if return_type_node else 'void'

                members.append({
                    'type': 'method',
                    'name': name,
                    'parameters': parameters,
                    'return_type': return_type
                })

        return members

    def _extract_parameters(self, node, source_code: str) -> List[Dict[str, Any]]:
        """提取函数或方法的参数"""
        parameters = []

        for child in node.children:
            # 跳过注释
            if child.type in ['comment', 'line_comment', 'block_comment']:
                continue

            # 改进点1：使用更可靠的方式获取参数名
            name = 'unknown'
            # 首先尝试通过字段名获取
            name_node = child.child_by_field_name('name')
            if name_node:
                name = self._get_node_text(name_node, source_code)
            else:
                # 如果通过字段名获取失败，尝试在子节点中查找标识符
                for subchild in child.children:
                    if subchild.type in ['identifier', 'property_identifier', 'type_identifier']:
                        name = self._get_node_text(subchild, source_code)
                        break

            # 改进点2：使用更可靠的方式获取参数类型
            param_type = 'any'
            type_node = child.child_by_field_name('type')
            if type_node:
                type_text = self._get_node_text(type_node, source_code)
                # 如果类型文本以冒号开头，去掉冒号
                if type_text.startswith(':'):
                    param_type = type_text[1:].strip()
                else:
                    param_type = type_text
            else:
                # 如果通过字段名获取失败，尝试在子节点中查找类型注解
                for subchild in child.children:
                    if subchild.type == 'type_annotation':
                        type_text = self._get_node_text(subchild, source_code)
                        if type_text.startswith(':'):
                            param_type = type_text[1:].strip()
                        else:
                            param_type = type_text
                        break

            # 确定参数是否可选
            is_optional = False
            if child.type == 'optional_parameter':
                is_optional = True
            else:
                # 检查是否有问号标记
                is_optional = any(c.type == 'question_mark' for c in child.children)

            # 根据参数类型处理
            if child.type == 'required_parameter':
                parameters.append({
                    'name': name,
                    'type': param_type,
                    'optional': is_optional
                })
            elif child.type == 'optional_parameter':
                parameters.append({
                    'name': name,
                    'type': param_type,
                    'optional': is_optional
                })

        return parameters

    def _extract_type_value(self, node, source_code: str) -> str:
        """提取类型别名的完整值"""
        # 使用查询来定位值节点
        query = self.language.query("""
        (type_alias_declaration
            value: (_) @type.value
        )
        """)

        captures = query.captures(node)
        for capture, _ in captures:
            if capture.type == 'type_alias_declaration':
                continue
            # 返回第一个匹配的值节点
            return self._get_node_text(capture, source_code)

        # 如果没有找到，使用备选方法
        cursor = node.walk()
        cursor.goto_first_child()

        # 跳过 name 节点
        if cursor.node.type == 'type_identifier' and cursor.node.field_name == 'name':
            cursor.goto_next_sibling()

        # 跳过等号节点
        if cursor.node.type == '=':
            cursor.goto_next_sibling()

        # 获取剩余部分作为值
        if cursor.node:
            return self._get_node_text(cursor.node, source_code)
        return 'unknown'

    def _extract_type_alias(self, node, source_code: str) -> Dict[str, Any]:
        """提取类型别名信息"""
        # 获取类型名
        name_node = node.child_by_field_name('name')
        name = self._get_node_text(name_node, source_code) if name_node else 'unknown'
    
        # 获取类型值
        value_node = node.child_by_field_name('value')
        value_text = self._get_node_text(value_node, source_code) if value_node else 'unknown'
    
        # 初始化参数列表和返回类型
        parameters = []
        return_type = None
    
        # 检查是否是函数类型
        if value_node and value_node.type == 'function_type':
            # 尝试获取参数列表
            params_node = value_node.child_by_field_name('parameters')
            if params_node:
                parameters = self._extract_parameters(params_node, source_code)
        
            # 尝试获取返回类型
            return_type_node = value_node.child_by_field_name('return_type')
            if return_type_node:
                return_type = self._get_node_text(return_type_node, source_code)
    
        return {
            'name': name,
            'value': value_text,
            'parameters': parameters,
            'return_type': return_type,
            'source_text': self._get_node_text(node, source_code)
        }

    def _extract_class(self, node, source_code: str) -> Dict[str, Any]:
        """提取类声明信息"""
        name_node = node.child_by_field_name('name')
        name = self._get_node_text(name_node, source_code) if name_node else 'unknown'

        return {
            'name': name,
            'source_text': self._get_node_text(node, source_code)
        }

    def _extract_variable(self, node, source_code: str) -> List[Dict[str, Any]]:
        """提取变量声明信息"""
        variables = []

        # 处理变量声明中的多个声明器
        for child in node.children:
            if child.type == 'variable_declarator':
                name_node = child.child_by_field_name('name')
                name = self._get_node_text(name_node, source_code) if name_node else 'unknown'

                # 尝试获取类型注解
                type_annotation = None
                if child.child_count > 0:
                    for grandchild in child.children:
                        if grandchild.type == 'type_annotation':
                            type_annotation = self._get_node_text(grandchild, source_code)
                            break

                variables.append({
                    'name': name,
                    'type': type_annotation or 'any',
                    'source_text': self._get_node_text(child, source_code)
                })

        return variables

    def _get_node_text(self, node, source_code: str) -> str:
        """获取节点的文本内容"""
        if not node:
            return ''
        start = node.start_byte
        end = node.end_byte
        return source_code[start:end]
    
    def print_syntax_tree(self, node, source_code: str, max_depth=None):
        """按树形关系打印语法树的所有节点信息"""
    
        print("=" * 80)
        print("语法树结构:")
        print("=" * 80)
    
        # 递归打印节点
        self._print_node(node, source_code, depth=0, max_depth=max_depth)

    def _print_node(self, node, source_code: str, depth=0, max_depth=None):
        """递归打印节点关键信息"""
        # 检查深度限制
        print(f'depth:{depth}, type: {node.type}')
    
        # 准备缩进
        indent = "  " * depth
    
        # 获取节点文本（截断过长的文本）
        node_text = self._get_node_text(node, source_code)
        print(f"{indent}{node.type}:\n {node_text}")

        for field_name in node.id:
            # 获取该字段名的子节点
            field_node = node.child_by_field_name(field_name)
            if field_node is not None:
                print(f"{indent}{field_name}:{self._get_node_text(field_node, source_code)}") 
    
        # 递归处理子节点
        for child in node.children:
            self._print_node(child, source_code, depth + 1, max_depth)

